# -*- coding: UTF-8 -*-
"""
@Project ：ai-demo
@File    ：asr_tts_pipeline_ordered.py
@IDE     ：PyCharm
@Author  ：zhangyiheng
@Date    ：2025/4/1 12:04
@Describe：ASR + TTS
"""
import os
import re
import io
import queue
import threading
from datetime import datetime
from pathlib import Path
from typing import Union

try:
    import sounddevice as sd
    import soundfile as sf
except ImportError:
    raise ImportError("需要安装 sounddevice 和 soundfile 库，请执行：pip install sounddevice soundfile")

from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
from modelscope.outputs import OutputKeys
from funasr import AutoModel

# 模型根目录配置
MODEL_ROOT = "/Users/zhangzhan/Desktop/zhangzhan/modelscopes/repository-c"


class OrderedAudioPlayer:
    """保序音频播放器"""

    def __init__(self):
        self.play_queue = queue.PriorityQueue()  # 优先级队列
        self.current_seq = 0
        self.play_thread = threading.Thread(target=self._play_worker, daemon=True)
        self.lock = threading.Lock()
        self.is_running = True
        self.condition = threading.Condition()  # 初始化condition
        self.play_thread.start()

    def _play_worker(self):
        """播放工作线程（严格按顺序）"""
        expected_seq = 0
        while self.is_running or not self.play_queue.empty():
            try:
                # 获取下一个应该播放的音频
                with self.condition:
                    while True:
                        try:
                            priority, (seq, audio_data) = self.play_queue.get_nowait()
                            if seq == expected_seq:
                                self._play_audio(audio_data)
                                expected_seq += 1
                                self.play_queue.task_done()
                                break
                            else:
                                # 不是期待的序号，重新放回队列
                                self.play_queue.put((priority, (seq, audio_data)))
                                self.condition.wait(0.1)  # 短暂等待
                        except queue.Empty:
                            if not self.is_running:
                                return
                            self.condition.wait(0.5)
                            break
            except Exception as e:
                print(f"播放出错: {str(e)}")

    def _play_audio(self, audio_data: bytes):
        """实际播放音频"""
        with io.BytesIO(audio_data) as f:
            data, sample_rate = sf.read(f)
            sd.play(data, sample_rate)
            sd.wait()

    def add_audio(self, audio_data: bytes, seq: int):
        """添加音频到播放队列"""
        with self.lock:
            self.play_queue.put((seq, (seq, audio_data)))
            with self.condition:
                self.condition.notify()

    def stop(self):
        """停止播放器"""
        self.is_running = False
        with self.condition:
            self.condition.notify_all()
        self.play_thread.join()


class SpeechProcessor:
    def __init__(self, max_workers: int = 4):
        """
        初始化语音处理器
        :param max_workers: 最大TTS工作线程数
        """
        # ASR模型配置
        self.asr_model = self._init_asr_model()
        # TTS模型配置
        self.tts_pipeline = self._init_tts_pipeline()
        # 保序播放器
        self.audio_player = OrderedAudioPlayer()
        # 线程池信号量
        self.semaphore = threading.Semaphore(max_workers)
        # 序列号计数器
        self.seq_counter = 0
        self.seq_lock = threading.Lock()

    def __del__(self):
        """清理资源"""
        self.audio_player.stop()

    def _init_asr_model(self):
        """初始化中文语音识别模型"""
        return AutoModel(
            model=f"{MODEL_ROOT}/iic/speech_paraformer-large-vad-punc_asr_nat-zh-cn-16k-common-vocab8404-pytorch",
            model_revision="v2.0.5",
            vad_model=f"{MODEL_ROOT}/iic/speech_fsmn_vad_zh-cn-16k-common-pytorch",
            vad_model_revision="v2.0.4",
            punc_model=f"{MODEL_ROOT}/iic/punc_ct-transformer_zh-cn-common-vocab272727-pytorch",
            punc_model_revision="v2.0.4",
            disable_update=True,
            force_download=False,
            resume_download=False,
            cache_dir=MODEL_ROOT,
        )

    def _init_tts_pipeline(self):
        """初始化语音合成管道"""
        return pipeline(
            task=Tasks.text_to_speech,
            model=f"{MODEL_ROOT}/iic/speech_sambert-hifigan_tts_zh-cn_16k",
            model_revision="v1.0.2"
        )

    def _get_next_seq(self) -> int:
        """获取下一个序列号"""
        with self.seq_lock:
            self.seq_counter += 1
            return self.seq_counter - 1  # 从0开始计数

    def speech_to_text(self, audio_path: str) -> str:
        """语音识别（同步）"""
        if not Path(audio_path).exists():
            raise FileNotFoundError(f"音频文件不存在: {audio_path}")
        result = self.asr_model.generate(input=audio_path)
        return self._extract_asr_text(result)

    def _extract_asr_text(self, result) -> str:
        """提取ASR结果中的文本"""
        if isinstance(result, list):
            return result[0].get("text", "").strip()
        elif isinstance(result, dict):
            return result.get("text", "").strip()
        return str(result).strip()

    def _tts_worker(self, text: str, seq: int):
        """TTS工作线程（带序列号）"""
        try:
            with self.semaphore:
                output = self.tts_pipeline(input=text, voice="zhiyan_emo")
                self.audio_player.add_audio(output[OutputKeys.OUTPUT_WAV], seq)
        except Exception as e:
            print(f"[句子#{seq}] TTS合成失败: {str(e)}")

    def text_to_speech(self, text: str, output_path: str = None,
                       voice: str = "zhiyan_emo", save_file: bool = True) -> Union[str, None]:
        """
        语音合成（异步保序版）
        :return: 保存文件时返回路径，否则返回None
        """
        if save_file:
            output = self.tts_pipeline(input=text, voice=voice)
            audio_data = output[OutputKeys.OUTPUT_WAV]

            if output_path is None:
                clean_text = ''.join(c for c in text[:20] if c.isalnum() or c in ('_', '-', ' '))
                clean_text = clean_text.replace(' ', '_').strip('_')
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                output_path = f"tts_{clean_text}_{timestamp}.wav"

            os.makedirs(os.path.dirname(output_path)) if os.path.dirname(output_path) else None
            with open(output_path, "wb") as f:
                f.write(audio_data)
            return os.path.abspath(output_path)
        else:
            seq = self._get_next_seq()
            t = threading.Thread(
                target=self._tts_worker,
                args=(text, seq),
                daemon=True
            )
            t.start()
            return None

    def split_text_into_sentences(self, text: str) -> list:
        """智能文本分段（改进版）"""
        # 处理中英文及各种标点
        sentences = re.split(r'(?<=[。！？.!?;；])', text)

        # 合并过短片段
        min_length = 8  # 最小字符数
        merged = []
        buffer = ""

        for s in [s.strip() for s in sentences if s.strip()]:
            if len(buffer) + len(s) < min_length:
                buffer += s
            else:
                if buffer:
                    merged.append(buffer)
                    buffer = s
                else:
                    merged.append(s)

        if buffer:
            merged.append(buffer)

        return merged

    def process_audio_pipeline(self, input_audio: str, output_audio: str = None,
                               save_to_file: bool = False) -> dict:
        """
        完整处理流水线（保序异步版）
        :param save_to_file: 为True时同步处理并保存文件
        :return: 包含asr_text和success字段的字典
        """
        results = {"asr_text": None, "success": False}

        try:
            # 语音识别（同步）
            results["asr_text"] = self.speech_to_text(input_audio)

            if results["asr_text"]:
                sentences = self.split_text_into_sentences(results["asr_text"])

                if save_to_file:
                    # 同步处理模式
                    tts_path = self.text_to_speech(
                        text=results["asr_text"],
                        output_path=output_audio,
                        save_file=True
                    )
                    results["tts_path"] = tts_path
                else:
                    # 异步保序处理模式
                    for sentence in sentences:
                        self.text_to_speech(
                            text=sentence,
                            save_file=False
                        )

                results["success"] = True

        except Exception as e:
            print(f"处理失败: {str(e)}")

        return results


if __name__ == "__main__":
    # 初始化处理器
    processor = SpeechProcessor(max_workers=3)  # 限制并发数

    # 测试案例
    test_audio = "/Users/zhangzhan/Desktop/zhangzhan/modelscopes/vad_example.wav"

    # 运行异步播放模式
    print("开始处理音频（异步模式）...")
    result = processor.process_audio_pipeline(
        input_audio=test_audio,
        save_to_file=False
    )

    # 打印识别结果
    print("\n" + "=" * 50)
    print(f"ASR识别结果: {result['asr_text']}")
    print("处理状态:", "成功" if result["success"] else "失败")

    # 保持程序运行
    input("按Enter键退出（后台可能仍在播放音频）...")